import os
import io
import torch
import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
import numpy as np
from config import data_transforms
from model import AIModel  # 导入模型类，根据实际情况调整
from video_processor import VideoProcessor
import shutil
from pathlib import Path

# 创建FastAPI应用
app = FastAPI(
    title="人脸检测API",
    description="提供人脸真伪检测服务的REST API",
    version="1.0.0"
)

# 添加CORS中间件，允许跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，实际生产环境应该限制
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 创建输出目录
UPLOAD_DIR = Path("upload_files")
OUTPUT_DIR = Path("output_videos")
UPLOAD_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)

# 挂载静态文件目录
app.mount("/videos", StaticFiles(directory=str(OUTPUT_DIR)), name="videos")

# 全局变量存储模型
model = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
video_processor = None

@app.on_event("startup")
async def load_model():
    """启动时加载模型"""
    global model, video_processor
    model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "best_model.pth")
    
    if not os.path.exists(model_path):
        raise RuntimeError(f"模型文件不存在: {model_path}")
    
    try:
        # 加载模型
        print(f"正在加载模型: {model_path}")
        model = torch.load(model_path, map_location=device)
        # 如果加载的是state_dict，则需要先实例化模型
        if isinstance(model, dict):
            model_instance = AIModel()  # 根据实际模型类进行调整
            model_instance.load_state_dict(model)
            model = model_instance
        model.to(device)
        model.eval()
        
        # 初始化视频处理器
        video_processor = VideoProcessor(model, device)
        
        print(f"模型加载成功，使用设备: {device}")
    except Exception as e:
        print(f"模型加载失败: {str(e)}")
        raise

async def predict_image(image_bytes):
    """
    对图像进行预测
    Args:
        image_bytes: 图像字节数据
    Returns:
        预测结果
    """
    try:
        # 打开图像
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
        # 应用预处理
        image_tensor = data_transforms["test"](image)
        image_tensor = image_tensor.unsqueeze(0).to(device)
        
        # 预测
        with torch.no_grad():
            output = model(image_tensor)
            prob = output.float().squeeze().cpu().numpy()
            
        # 返回预测结果
        prediction = 1 if prob >= 0.5 else 0
        return {
            "success": True,
            "prediction": int(prediction),  # 1表示伪造，0表示真实
            "probability": float(prob),
            "label": "伪造" if prediction == 1 else "真实",
            "confidence": float(prob) if prediction == 1 else float(1 - prob)
        }
    except Exception as e:
        return {"success": False, "error": str(e)}

def cleanup_temp_file(file_path):
    """清理临时文件"""
    try:
        if os.path.exists(file_path):
            os.unlink(file_path)
    except Exception as e:
        print(f"清理临时文件失败: {str(e)}")

@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    """
    检测上传图像是真实人脸还是伪造人脸
    
    - **file**: 要检测的图像文件
    
    返回:
    - **prediction**: 0表示真实，1表示伪造
    - **probability**: 模型输出的原始概率值
    - **label**: 文本标签（真实/伪造）
    - **confidence**: 置信度
    """
    # 检查文件类型
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="只接受图像文件")
    
    # 读取图像内容
    contents = await file.read()
    
    # 预测
    result = await predict_image(contents)
    
    if not result["success"]:
        raise HTTPException(status_code=500, detail=result["error"])
    
    return JSONResponse(content=result)

@app.post("/detect_video")
async def detect_video(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    sample_rate: int = Form(10),
    max_frames: int = Form(300)
):
    """
    检测视频中的人脸真伪
    
    - **file**: 要检测的视频文件
    - **sample_rate**: 采样率（每隔多少帧处理一次），默认10
    - **max_frames**: 最大处理帧数，默认300
    
    返回:
    - **overall_prediction**: 总体预测结果（0真实，1伪造）
    - **fake_ratio**: 伪造帧比例
    - **frames_processed**: 处理的帧数
    - **video_url**: 带标记的输出视频URL
    """
    # 检查文件类型
    if not file.content_type.startswith("video/"):
        raise HTTPException(status_code=400, detail="只接受视频文件")
    
    # 检查是否加载了模型
    if video_processor is None:
        raise HTTPException(status_code=500, detail="视频处理器未初始化")
    
    # 读取视频内容
    contents = await file.read()
    
    try:
        # 处理视频
        results, output_path = video_processor.process_video(
            contents, 
            sample_rate=sample_rate,
            max_frames=max_frames,
            output_video=True
        )
        
        # 移动输出视频到静态目录
        if output_path and os.path.exists(output_path):
            output_filename = os.path.basename(output_path)
            destination = OUTPUT_DIR / output_filename
            shutil.move(output_path, destination)
            
            # 设置视频URL
            video_url = f"/videos/{output_filename}"
            results["video_url"] = video_url
            
            # 添加后台任务清理输出视频（可选，根据需要调整）
            # background_tasks.add_task(cleanup_temp_file, str(destination))
        
        return JSONResponse(content=results)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"视频处理失败: {str(e)}")

@app.get("/videos/{filename}")
async def get_video(filename: str):
    """获取处理后的视频文件"""
    video_path = OUTPUT_DIR / filename
    if not video_path.exists():
        raise HTTPException(status_code=404, detail="视频文件不存在")
    return FileResponse(str(video_path))

@app.get("/health")
async def health_check():
    """健康检查接口"""
    return {"status": "healthy", "model_loaded": model is not None}

if __name__ == "__main__":
    # 运行服务器
    uvicorn.run("api:app", host="0.0.0.0", port=5000, reload=True) 